package com.ruoyi.system.websocket;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.ruoyi.common.utils.DateUtils;
import com.ruoyi.system.websocket.body.MessageBody;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket服务类
 */
@Component
@Slf4j
@ServerEndpoint("/traffic")
public class WebSocketServer {
    /**
     * 心跳消息
     */
    private final static String PING = "ping";

    private final static String PONG = "pong";

    /**
     * 存放每个客户端对应的 WebSocketServer 对象
     */
    public static ConcurrentHashMap<String, WebSocketServer> webSocketMap = new ConcurrentHashMap<>();

    /**
     * 与某个客户端的连接会话，需要通过它来给客户端发送数据
     */
    private Session session;

    /**
     * 接收 userId
     */
    private String userId = "1";

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        // 加入
        webSocketMap.put(userId, this);
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        if (!webSocketMap.containsKey(userId)) {
            return;
        }
        webSocketMap.remove(userId);
    }

    /**
     * 收到客户端消息后调用的方法
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        MessageBody messageBody = null;
        try {
            // 解析消息体
            messageBody = JSONObject.parseObject(message, MessageBody.class);
        } catch (Exception e) {
            // websocket消息内容异常
//            throw new BusinessException("A-030002");
        }
        // 消息内容为PING,返回PONG给客户端
        if ("HEART_BEAT".equals(messageBody.getMsgType()) && PING.equals(messageBody.getMsgContent())) {
            try {
                MessageBody pongMessageBody = new MessageBody();
                pongMessageBody.setMsgContent(PONG);
                pongMessageBody.setMsgType("HEART_BEAT");
                pongMessageBody.setTime(DateUtils.getTime());
                pongMessageBody.setClientMsgId(String.valueOf(System.currentTimeMillis()));
                pongMessageBody.setServerMsgId(messageBody.getClientMsgId());
                pongMessageBody.setUserId(messageBody.getUserId());
                // 响应PONG消息给客户端
                this.sendMessage(session, JSON.toJSONString(pongMessageBody));
            } catch (Exception e) {
                log.error("发送消息出错", e);
            }
        }
    }

    /**
     * 发生错误时调用
     */
    @OnError
    public void onError(Session session, Throwable error) {
        error.printStackTrace();
    }

    /**
     * 实现服务器主动推送
     */
    public void sendMessage(Session session, String message) throws IOException {
        session.getBasicRemote().sendText(message);
        log.info("消息推送到前端成功:", message);
    }

    /**
     * 群发消息
     */
    public void sendMessageToAll(String message) throws IOException {
        for (Map.Entry<String, WebSocketServer> entry : webSocketMap.entrySet()) {
            WebSocketServer webSocketServer = entry.getValue();
            webSocketServer.sendMessage(this.session,message);
        }
    }

    /**
     * 单发消息
     */
    public static void sendMessageToUser(String toUserId, String message) throws IOException, EncodeException {
        if (webSocketMap.containsKey(toUserId)) {
            WebSocketServer webSocketServer = webSocketMap.get(toUserId);
            webSocketServer.sendMessage(webSocketServer.session, message);
        } else {
        }
    }

    /**
     * 获取在线人数
     */
    public static int getOnlineCount() {
        return webSocketMap.size();
    }

    /**
     * 用户是否在线
     */
    public static Boolean isOnline(String userId) {
        return webSocketMap.containsKey(userId);
    }

    /**
     * 在线用户
     */
    public static Set<String> getOnlineUsers() {
        Set<String> set = new HashSet<>();
        Enumeration<String> enumeration = webSocketMap.keys();

        while (enumeration.hasMoreElements()) {
            set.add(enumeration.nextElement());
        }

        return set;
    }
}

